import numpy
import math
import common.states


class Hamiltonian_N_sz:
    def __init__(self, **kwargs):
        self.kMatrix = kwargs['kMatrix']
        self.nElectrons = kwargs['nElectrons']
        self.states = kwargs['states']
        self.nStates = len(self.states)
        self.u = kwargs['u']
        self.conTable = common.states.conversionTable(self.states)
        hBlock = self.updateU()
        hBlock = self.updateT(hBlock)
        hBlock = self.updateMu(hBlock)
        self.h = hBlock


    def updateU(self):
        hBlock = numpy.matrix(self.nStates)
        for state in self.states:
            stateIndex = common.states.state2I(state, self.conTable)
            temp = 0
            for s in range(self.nSites):
                temp += (state.up[s] - 0.5) * (state.down[s] - 0.5)
            hBlock[stateIndex][stateIndex] = self.u * temp;

        return hBlock

    def updateT(self, hBlock_):
        hBlock = hBlock_.clone()
        for state in self.states:
            stateIndex = common.states.state2I(state, self.conTable)
            for i in range(self.nSites):
                for j in range(self.nSites):
                    if i == j:
                        continue
                    t_up = self.kMatrix.up[i][j]
                    if t_up != 0:
                        hBlock = self.cCrossC('up', i, j, state, stateIndex, t_up, hBlock)

                    t_down = self.kMatrix.down[i][j]
                    if t_up != 0:
                        hBlock = self.cCrossC('down', i, j, state, stateIndex, t_down, hBlock)

        return hBlock

    def updateMu(self, hBlock_):
        hBlock = hBlock_.clone()
        for i in range(self.nStates):
            hBlock[i][i] -= self.mu * self.nElectrons;
        return hBlock

    def cCrossC(self, electronType, i, j, state, stateIndex, t, hBlock):
        tempState = state.clone()
        if electronType == 'up':
            if tempState.up[j] == 1 and tempState.up[i] == 0:
                minBoundary = min(i, j)
                maxBoundary = max(i, j)
                factor = 1
                if minBoundary + 1 != maxBoundary:
                    factor = math.pow(-1, sum(state.up, minBoundary + 1, maxBoundary - 1))

                result = hBlock.clone()

                tempState.lowerUp(j)
                tempState.raiseUp(i)

                tempStateIndex = common.states.state2I(tempState, self.conTable)
                hBlock[stateIndex][tempStateIndex] = -t * factor
                hBlock[tempStateIndex][stateIndex] = -t * factor

                return result

        if electronType == 'down':
            if tempState.down[j] == 1 and tempState.down[i] == 0:
                minBoundary = min(i, j)
                maxBoundary = max(i, j)
                factor = 1
                if minBoundary + 1 != maxBoundary:
                    factor = math.pow(-1, sum(state.up, minBoundary + 1, maxBoundary - 1))

                result = hBlock.clone()

                tempState.lowerDown(j)
                tempState.raiseDown(i)

                tempStateIndex = common.states.state2I(tempState, self.conTable)
                hBlock[stateIndex][tempStateIndex] = -t * factor
                hBlock[tempStateIndex][stateIndex] = -t * factor

                return result;
        return hBlock